####此代码用于对基因组基因所有CRISPR-Cas9系统识别的靶点(NNNNNNNNNNNNNNNNNNNNNGG)进行预测与统计####
####此代码由陈大嵩编辑，邮箱chendasong@gmail.com####

from Bio import SeqIO
import os
import argparse
from collections import Counter

parser = argparse.ArgumentParser(description='利用基因组与GFF注释信息预测')
parser.add_argument('-f', '--fasta', type=str, help='基因组fasta文件')
parser.add_argument('-g', '--gff', type=str, help='基因组注释GFF文件')
parser.add_argument('-o', '--out', type=str, default='CRISPR', help='输出文件')
parser.add_argument('-t', '--threads', type=int, default=12, help='使用CPU核心数')
parser.add_argument('-m', '--mismatches', type=int, default=1, help='脱靶位点数（1-3）')
args = parser.parse_args()
genome_file = args.fasta
gff_file = args.gff
out_file = args.out
threads = args.threads
v = args.mismatches
if not 1<= v <= 3:
    print('脱靶位点数只能选1、2、3')
    exit()

complement_seq = str.maketrans('ATGC','TACG') #生成互补序列

print('读取基因组序列')
def genome_to_dic(genome):
    return SeqIO.to_dict(SeqIO.parse(genome, 'fasta'))
genome = {}
try:
    genome = genome_to_dic(genome_file) #读取基因组序列到字典    
except:
    print('基因组文件输入错误')
    exit()

print('预测非重复CRISPR位点')
def crispr_perdict(fastafile): #预测基因组中CRISPR的序列与数量，并存于字典中
    crispr_list = []
    #n = 0
    for seq in fastafile:
        forword_seq = str(fastafile[seq].seq).upper()
        reversed_seq = str(fastafile[seq].reverse_complement().seq).upper()
        for i in range(21,len(forword_seq)):
            if forword_seq[i:i+2] == 'GG':
                crispr_list.append(forword_seq[i-21:i-1])
        for i in range(21,len(reversed_seq)):
            if reversed_seq[i:i+2] == 'GG':
                crispr_list.append(reversed_seq[i-21:i-1])
    crispr_dic = dict(Counter(crispr_list))
    crispr_dic_one = {k:i for k, i in crispr_dic.items() if i == 1 and k.count('N') == 0}
    return crispr_dic_one
crispr_one = crispr_perdict(genome) #提取所有基因组上的非重复CRISPR位点


print('提取基因的CRISPR位点')
def gene_iso(gff_file): #提取基因序列并预测非重复CRISPR位点
    gene_predic = {}
    gff_list = []
    for line in open(gff_file):
        gff_list.append(line.strip())
        if not line[0] == '#':
            line = line.split('\t')
            if len(line) == 9 and line[2] == 'gene':
                gene_id = line[-1].split(';')[0].split('-')[-1]
                genome_id = line[0]
                genome_seq = str(genome[line[0]].seq.upper())
                start = int(line[3])
                end = int(line[4])
                for i in range(start-1,end-22):
                    f = genome_seq[i:i+23]
                    r = f[::-1].translate(complement_seq)
                    if f[-2:] == 'GG' and f[:-3] in crispr_one:
                        gene_predic[f] = [gene_id, genome_id, i+1, i+23, '+']
                    if r[-2:] == 'GG' and r[:-3] in crispr_one:
                        gene_predic[r] = [gene_id, genome_id, i+1, i+23, '-']
    return gene_predic
crispr_gene = {}
try:
    CRISPR_gene = gene_iso(gff_file)
except:
    print('GFF文件输入错误')
    exit()


print('输出CRISPR位点')
with open(out_file+'.fas', 'w') as res:
    for c in CRISPR_gene:
        print(''.join(['>',CRISPR_gene[c][0],'|',CRISPR_gene[c][1],':',str(CRISPR_gene[c][2]),'-',str(CRISPR_gene[c][3]),'(',CRISPR_gene[c][4],')']), c, file = res, sep = '\n')

print('预测脱靶位点')
os.system(' '.join(['bowtie-build -q --threads', str(threads), genome_file, out_file]))
os.system(' '.join(['bowtie --sam-nohead --trim3 3 -a -f -v', str(v), '-p', str(threads), out_file, out_file+'.fas', '-S',out_file+'.sam']))

print('统计脱靶位点')
def sam_stat(sam_file):
    sam_dic = {}
    for line in open(sam_file):
        line = line.split('\t')    
        if not line[0] in sam_dic:
            sam_dic[line[0]] = [0,0,0,0] #储存脱靶0、1、2、3个位点的碱基
        base1 = int(line[3])-1
        if genome[line[2]].seq[base1-3:base1-1] == 'CC' or genome[line[2]].seq[base1+21:base1+23] == 'GG': #判断比对位点是否可以被Cas9识别
            n = int(line[-4].split(':')[-1])
            try:
                sam_dic[line[0]][n] += 1
            except:
                print('\t'.join(line))
    return sam_dic
sam_stat = sam_stat(out_file+'.sam')
with open(out_file+'_stat.txt', 'w') as sam_stat_res:
    print('Target', '1 off_target' , '2 off_target', '3 off_target', file = sam_stat_res, sep = '\t')
    for i in sam_stat:
        print(i, sam_stat[i][1], sam_stat[i][2], sam_stat[i][3], file = sam_stat_res, sep = '\t')

print('Done !!!')